import networkx as nx
from collections import deque
from pulp import *
from tools import generate_gas_network, visualize_network, get_outward_connected_components
from matplotlib.font_manager import FontProperties
font = FontProperties(fname='resources/SimHei.ttf', size=12)  # 替换为实际字体路径

def calculate_users(H, sources, leak_impact_edges):
    G = H.copy()
    for u, v in leak_impact_edges:
        G.remove_edge(u, v)
    users = 0
    for source in sources:
        queue = deque([source])
    visited = set([source])
    while queue:
        node = queue.popleft()
        if G.nodes[node].get("user_weight"):
            users += G.nodes[node]["user_weight"]
        for neighbor in G.successors(node):
            if neighbor not in visited:
                visited.add(neighbor)
                queue.append(neighbor)
    return users

def calculate_leaks(G, sources):
    leaks = 0
    for source in sources:
        queue = deque([source])
    visited = set([source])
    while queue:
        node = queue.popleft()
        if G.nodes[node].get("leak_weight"):
            leaks += G.nodes[node]["leak_weight"]
        for neighbor in G.successors(node):
            if neighbor not in visited:
                visited.add(neighbor)
                queue.append(neighbor)
    return leaks

def calculate_impact_users(G, valves_to_cut, leak_impact_edges):
    """计算关闭阀门的影响范围（用户权重和）"""
    H = G.copy()
    for u, v in G.edges():
        if u in valves_to_cut:
            H.remove_edge(u, v)
    sources_G = [n for n, attr in G.nodes(data=True) if attr['type'] == 'source']
    sources_H = [n for n, attr in H.nodes(data=True) if attr['type'] == 'source']            
    return calculate_users(G, sources_G, leak_impact_edges) - calculate_users(H, sources_H, leak_impact_edges)

def calculate_impact_leaks(G, valves_to_cut):
    """计算关闭阀门的影响范围（泄漏权重和）"""
    H = G.copy()
    for u, v in G.edges():
        if u in valves_to_cut:
            H.remove_edge(u, v)
    sources_H = [n for n, attr in H.nodes(data=True) if attr['type'] == 'source']   
    return calculate_leaks(H, sources_H)

def generate_super_source_G(G, leak_impact_edges):
    leaks = [n for n, attr in G.nodes(data=True) if attr['type'] == 'leak']
    sources = [n for n, attr in G.nodes(data=True) if attr['type'] == 'source']        
    H = G.copy()
    super_source = "super_source"
    super_sink = "super_sink"
    nodes = []
    nodes.append((super_source, {'type': 'super_source'}))
    nodes.append((super_sink, {'type': 'super_sink'}))
    H.add_nodes_from(nodes)

    for u in sources:
        H.add_edge(super_source, u, capacity=1e9)
    for u in leaks:
        H.add_edge(u, super_sink, capacity=1e9)

    for u, v in H.edges:
        if u in valves:
            H[u][v]['capacity'] = 1
        else:
            H[u][v]['capacity'] = 1e9  
    for u, v in leak_impact_edges:
        H[u][v]['capacity'] = 0
            
    return H, super_source, super_sink

def generate_super_source_users_G(G, uv_dict, leak_impact_edges):
    leaks = [n for n, attr in G.nodes(data=True) if attr['type'] == 'leak']
    sources = [n for n, attr in G.nodes(data=True) if attr['type'] == 'source']    
    H = G.copy()
    super_source = "super_source"
    super_sink = "super_sink"
    nodes = []
    nodes.append((super_source, {'type': 'super_source'}))
    nodes.append((super_sink, {'type': 'super_sink'}))
    H.add_nodes_from(nodes)

    for u in sources:
        H.add_edge(super_source, u, capacity=1e9)
    for u in leaks:
        H.add_edge(u, super_sink, capacity=1e9)

    for u, v in H.edges:
        if u in valves:
            H[u][v]['capacity'] = uv_dict[(u,v)]
        else:
            H[u][v]['capacity'] = 1e9
    for u, v in leak_impact_edges:
        H[u][v]['capacity'] = 0            
    return H, super_source, super_sink

def find_min_cut(G, H, super_source, super_sink, valves):
    """基于流向的有向图最小割算法"""
    # 计算最小割
    min_cut_valve, partition = nx.minimum_cut(H, super_source, super_sink)
    print ("partition", partition)
    print ("min_cut_valve", min_cut_valve)
    reachable, non_reachable = partition

    min_cut_set = []
    a = 0
    for u, v in H.edges():
        if u in reachable and v in non_reachable:
            if G.has_edge(u, v) and u in valves:
                min_cut_set.append((u, v))
                a += H[u][v]['capacity']
    return min_cut_set

def get_leak_impact_edges(G, sources, leaks):
    H = G.copy()
    for u, v in G.edges():
        if u in leaks:
            H.remove_edge(u, v)
    _, reachable_edges = get_outward_connected_components(H, sources)
    edges = []
    for u, v in G.edges():
        edges.append((u,v))
    no_reachable_edges = set(edges) - reachable_edges
    return no_reachable_edges

def generate_cost_flow_network(H, leaks):
    G = H.copy()
    for u, v in H.edges():
        if u in leaks:
            G.remove_edge(u, v)
    
    
    
if __name__ == "__main__":
    G_ori, pos = generate_gas_network()
    leaks = [n for n, attr in G_ori.nodes(data=True) if attr['type'] == 'leak']
    sources = [n for n, attr in G_ori.nodes(data=True) if attr['type'] == 'source']
    valves = [n for n, attr in G_ori.nodes(data=True) if attr['type'] == 'valve']
    users = [n for n, attr in G_ori.nodes(data=True) if attr['type'] == 'user']
    print ("leaks", leaks)
    #exit()
    leak_impact_edges = get_leak_impact_edges(G_ori, sources, leaks)
    print (leak_impact_edges)
    #exit()
    visualize_network(G_ori, pos, "燃气管网可视化", leak_impact_edges=leak_impact_edges)
    #exit()
    print ("源头数:", len(sources))
    print ("阀门数:", len(valves))
    print ("泄漏数:", len(leaks))
    print ("用户总数:", len(users))

    # # 直接最小割
    # print ("#####################")
    # print ("------直接最小割------")
    # G = G_ori.copy()
    # H, super_source, super_sink = generate_super_source_G(G, leak_impact_edges)
    # min_cut_set = find_min_cut(G, H, super_source, super_sink, valves)
    # valves_to_cut = []
    # for u, _ in min_cut_set:
    #     valves_to_cut.append(u)
    # print(f"定位到关键阀门: {valves_to_cut}")
    # print("阀门影响用户数:", calculate_impact_users(G, valves_to_cut, leak_impact_edges))
    # print("未阻止泄漏数:", calculate_impact_leaks(G, valves_to_cut))
    # visualize_network(H, pos, "燃气管网可视化——关阀1", min_cut_set, leak_impact_edges)

    # 损失用户最小最小割
    print ("###########################")
    print ("------损失用户最小最小割------")
    G = G_ori.copy()
    weight_dict, capacity_dict = calculate_dict(G)
    # for valve in valves:
    #     num = calculate_users(G, [valve], leak_impact_edges)
    #     for v in G.successors(valve):
    #         uv_dict[(valve,v)] = num
    #         print (valve, v, num)
    H, super_source, super_sink = generate_super_source_users_G(G, uv_dict, leak_impact_edges)
    # maxflow = nx.maximum_flow(H, super_source, super_sink)[1]
    # print ("maxflow", maxflow)
    # uv_dict = {}
    # for valve, vv in maxflow.items():
    #     if valve in valves:
    #         for v, num in vv.items():
    #             if num > 0:
    #                 uv_dict[(valve,v)] = num
    #H, super_source, super_sink = generate_super_source_users_G(G, uv_dict)
    min_cut_set = find_min_cut(G, H, super_source, super_sink, valves)
    valves_to_cut = []
    for u, _ in min_cut_set:
        valves_to_cut.append(u)
    print(f"定位到关键阀门: {valves_to_cut}")
    print("阀门影响用户数:", calculate_impact_users(G, valves_to_cut, leak_impact_edges))
    print("未阻止泄漏数:", calculate_impact_leaks(G, valves_to_cut)) 
    for (u, v), num in uv_dict.items():
        H[u][v]['capacity'] = num
    maxflow = nx.maximum_flow(H, super_source, super_sink)[1]
    print ("maxflow", maxflow)
    uv_dict = {}
    for valve, vv in maxflow.items():
        if valve in valves or True:
            for v, num in vv.items():
                if num > 0 or True:
                    H[valve][v]['capacity'] = num
    visualize_network(H, pos, "燃气管网可视化——关阀2", min_cut_set, leak_impact_edges)   
    